from Crypto.Cipher import AES, DES3
from PyQt5.QtCore import QObject, pyqtSignal
import os


class Cipher(QObject):
    # cipher信号，当加（解）密操作开时发送此信号
    cipher_start = pyqtSignal(str, int)
    # cipher信号，在加（解）密过程中此信号将不停地被发送
    cipher_variety = pyqtSignal(str, int)

    def __init__(self, **kwargs):
        super(Cipher, self).__init__()
        self.__infile = kwargs["infile"]
        self.__outdir = kwargs["outdir"]
        # 将密码通过ascii编码转化为二进制字符串
        self.__psw = kwargs["psw"].encode("ascii")
        alg = kwargs["alg"]
        # iv是AES、DES算法初始化向量，长度为16个字节
        iv = b"__yosuga_no_sora"
        # 通过从combobox中选择的算法名称，从pycrypto库中创建对应算法的对象
        # 同时根据特定的算法，获得特定算法的block_size存储在self.__block_size中
        if alg[0:3] == "AES":
            self.__cipher = AES.new(self.__psw, AES.MODE_CBC, iv)
            self.__block_size = AES.block_size
        else:
            self.__cipher = DES3.new(self.__psw, DES3.MODE_CBC, iv)
            self.__block_size = DES3.block_size

    def encrypt(self):
        # 当前文件不是‘.moe’后缀的文件，否则断言失败，程序异常
        assert not self.__infile.endswith(".moe")
        # 当前文件小于500MB，否则断言失败，程序异常
        assert os.path.getsize(self.__infile) < 536870912
        # 获取当前文件block_size的总数
        count = os.path.getsize(self.__infile) // self.__block_size + 1
        indir, infile = os.path.split(self.__infile)
        # 加密后‘.moe’文件存放位置的绝对路径
        outfile = self.__outdir + "/" + infile + ".moe"
        # 加密开始，发送cipher_start信号
        self.cipher_start.emit("encrypt", count)
        with open(outfile, "wb") as f_out:
            with open(self.__infile, "rb") as f_in:
                for i in range(count):
                    buf_0 = f_in.read(self.__block_size).ljust(self.__block_size, b'\x00')
                    f_out.write(self.__cipher.encrypt(buf_0))
                    # 向目标文件写入第i次已加密的block_size，发送cipher_variety信号
                    self.cipher_variety.emit("encrypt", i)
        # 加密完成，发送参数为-1的cipher_variety信号
        self.cipher_variety.emit("encrypt", -1)

    def decrypt(self):
        # 当前文件是‘.moe’后缀的文件，否则断言失败，程序异常
        assert self.__infile.endswith(".moe")
        # 获取当前文件block_size的总数
        count = os.path.getsize(self.__infile) // self.__block_size + 1
        indir, infile = os.path.split(self.__infile)
        # 解密后文件存放位置的绝对路径
        outfile = self.__outdir + "/" + infile[:-4]
        # 解密开始，发送cipher_start信号
        self.cipher_start.emit("decrypt", count)
        with open(outfile, "wb") as f_out:
            with open(self.__infile, "rb") as f_in:
                for i in range(count):
                    buf_0 = f_in.read(self.__block_size).ljust(self.__block_size, b'\x00')
                    f_out.write(self.__cipher.decrypt(buf_0))
                    # 向目标文件写入第i次已解密的block_size，发送cipher_variety信号
                    self.cipher_variety.emit("decrypt", i)
        # 解密完成，发送参数为-1的cipher_variety信号
        self.cipher_variety.emit("decrypt", -1)
